import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from neo4j import GraphDatabase

'''
此模块提供了一个函数，用于连接Neo4j数据库中的文档节点，
并根据余弦相似度创建相似文档之间的关系。

'''
def connect_similar_documents(uri: str = "bolt://localhost:7687", user: str = "neo4j", password: str = "12345678", similarity_threshold: float = 0.7):
    """连接Neo4j中相似的文档节点"""

    # 连接数据库
    driver = GraphDatabase.driver(uri, auth=(user, password))

    try:
        # 获取文档数据
        with driver.session() as session:
            result = session.run("MATCH (d:Document) RETURN d.id as id, d.embedding as embedding")
            doc_data = [(record["id"], record["embedding"]) for record in result]

        if len(doc_data) < 2:
            return []

        # 提取ID和嵌入
        doc_ids = [doc[0] for doc in doc_data]
        embeddings = np.array([doc[1] for doc in doc_data])

        # 计算相似度矩阵
        similarity_matrix = cosine_similarity(embeddings)

        # 找到相似文档对
        connections = []
        for i in range(len(doc_ids)):
            for j in range(i + 1, len(doc_ids)):
                similarity = similarity_matrix[i][j]
                if similarity >= similarity_threshold:
                    connections.append((doc_ids[i], doc_ids[j], float(similarity)))

        # 创建Neo4j关系（双向）
        with driver.session() as session:
            # 清除现有关系
            session.run("MATCH ()-[r:SIMILAR_TO]->() DELETE r")

            # 创建双向关系
            for doc1_id, doc2_id, similarity in connections:
                session.run("""
                    MATCH (d1:Document {id: $doc1_id})
                    MATCH (d2:Document {id: $doc2_id})
                    MERGE (d1)-[:SIMILAR_TO {similarity: $similarity}]->(d2)
                    MERGE (d2)-[:SIMILAR_TO {similarity: $similarity}]->(d1)
                    """,
                            doc1_id=doc1_id,
                            doc2_id=doc2_id,
                            similarity=similarity
                            )

        # 打印连接信息
        print(f"创建了 {len(connections)} 个连接")
        for doc1_id, doc2_id, similarity in connections[:10]:
            print(f"{doc1_id} <-{similarity:.3f}-> {doc2_id}")

        return connections

    finally:
        driver.close()


# 使用示例
if __name__ == "__main__":
    connect_similar_documents()